{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing packages:\n",
      "\t.package(url: \"https://github.com/JacopoMangiavacchi/SwiftCoreMLTools.git\", from: \"0.0.6\")\n",
      "\t\tSwiftCoreMLTools\n",
      "\t.package(url: \"https://github.com/dduan/Just.git\", from: \"0.8.0\")\n",
      "\t\tJust\n",
      "With SwiftPM flags: ['-c', 'release']\n",
      "Working in: /tmp/tmpdov4rwva/swift-install\n",
      "Fetching https://github.com/JacopoMangiavacchi/SwiftCoreMLTools.git\n",
      "Fetching https://github.com/dduan/Just.git\n",
      "Fetching https://github.com/apple/swift-protobuf.git\n",
      "Cloning https://github.com/dduan/Just.git\n",
      "Resolving https://github.com/dduan/Just.git at 0.8.0\n",
      "Cloning https://github.com/apple/swift-protobuf.git\n",
      "Resolving https://github.com/apple/swift-protobuf.git at 1.8.0\n",
      "Cloning https://github.com/JacopoMangiavacchi/SwiftCoreMLTools.git\n",
      "Resolving https://github.com/JacopoMangiavacchi/SwiftCoreMLTools.git at 0.0.6\n",
      "[1/3] Compiling Just Just.swift\n",
      "[2/3] Compiling SwiftProtobuf AnyMessageStorage.swift\n",
      "[3/4] Compiling SwiftCoreMLTools Activations.swift\n",
      "[4/5] Compiling jupyterInstalledPackages jupyterInstalledPackages.swift\n",
      "[5/5] Linking libjupyterInstalledPackages.so\n",
      "Initializing Swift...\n",
      "Installation complete!\n"
     ]
    }
   ],
   "source": [
    "%install-swiftpm-flags -c release\n",
    "%install '.package(url: \"https://github.com/JacopoMangiavacchi/SwiftCoreMLTools.git\", from: \"0.0.6\")' SwiftCoreMLTools\n",
    "%install '.package(url: \"https://github.com/dduan/Just.git\", from: \"0.8.0\")' Just"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TO4FT9bUohHx"
   },
   "outputs": [],
   "source": [
    "import Foundation\n",
    "import SwiftCoreMLTools\n",
    "import TensorFlow\n",
    "import Just"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Download\n",
    "\n",
    "Boston house prices dataset\n",
    "---------------------------\n",
    "\n",
    "**Data Set Characteristics:**  \n",
    "\n",
    "    :Number of Instances: 506 \n",
    "\n",
    "    :Number of Attributes: 13 numeric/categorical predictive. Median Value (attribute 14) is usually the target.\n",
    "\n",
    "    :Attribute Information (in order):\n",
    "        - CRIM     per capita crime rate by town\n",
    "        - ZN       proportion of residential land zoned for lots over 25,000 sq.ft.\n",
    "        - INDUS    proportion of non-retail business acres per town\n",
    "        - CHAS     Charles River dummy variable (= 1 if tract bounds river; 0 otherwise)\n",
    "        - NOX      nitric oxides concentration (parts per 10 million)\n",
    "        - RM       average number of rooms per dwelling\n",
    "        - AGE      proportion of owner-occupied units built prior to 1940\n",
    "        - DIS      weighted distances to five Boston employment centres\n",
    "        - RAD      index of accessibility to radial highways\n",
    "        - TAX      full-value property-tax rate per ten thousand dollars\n",
    "        - PTRATIO  pupil-teacher ratio by town\n",
    "        - B        1000(Bk - 0.63)^2 where Bk is the proportion of blacks by town\n",
    "        - LSTAT    % lower status of the population\n",
    "        - MEDV     Median value of owner-occupied homes in a thousand dollar\n",
    "\n",
    "    :Missing Attribute Values: None\n",
    "\n",
    "    :Creator: Harrison, D. and Rubinfeld, D.L.\n",
    "\n",
    "This is a copy of UCI ML housing dataset.\n",
    "https://archive.ics.uci.edu/ml/machine-learning-databases/housing/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "// if let cts = Just.get(URL(string: \"https://archive.ics.uci.edu/ml/machine-learning-databases/housing/housing.data\")!).content {\n",
    "//     try! cts.write(to: URL(fileURLWithPath:\"../data/housing.csv\"))\n",
    "// }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Ingestion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "let data = try String(contentsOfFile:\"../data/housing.csv\", encoding: String.Encoding.utf8)\n",
    "let dataRecords: [[Float]] = data.split(separator: \"\\n\").map{ String($0).split(separator: \" \").compactMap{ Float(String($0)) } }\n",
    "\n",
    "let numRecords = dataRecords.count\n",
    "let numColumns = dataRecords[0].count\n",
    "\n",
    "var index = Set<Int>()\n",
    "\n",
    "while index.count < numRecords {\n",
    "    index.insert(Int.random(in: 0..<numRecords))\n",
    "}\n",
    "\n",
    "let randomDataRecords = index.map{ dataRecords[$0] }\n",
    "\n",
    "let dataFeatures = randomDataRecords.map{ Array($0[0..<numColumns-1]) }\n",
    "let dataLabels = randomDataRecords.map{ Array($0[(numColumns-1)...]) }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Transformation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split Numerical Categorical Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "let categoricalColumns = [3, 8]\n",
    "let numericalColumns = [0, 1, 2, 4, 5, 6, 7, 9, 10, 11, 12]\n",
    "let numCategoricalFeatures = categoricalColumns.count\n",
    "let numNumericalFeatures = numericalColumns.count\n",
    "let numLabels = 1\n",
    "\n",
    "assert(numColumns == numCategoricalFeatures + numNumericalFeatures + 1)\n",
    "\n",
    "// Get Categorical Features\n",
    "let allCategoriesValues = dataFeatures.map{ row in categoricalColumns.map{ Int32(row[$0]) } }\n",
    "                                .reduce(into: Array(repeating: [Int32](), count: 2)){ total, value in\n",
    "                                    total[0].append(value[0])\n",
    "                                    total[1].append(value[1]) }\n",
    "                                .map{ Set($0).sorted() }\n",
    "\n",
    "let categoricalFeatures = dataFeatures.map{ row in categoricalColumns.map{ Int32(row[$0]) } }\n",
    "\n",
    "// Get Numerical Features\n",
    "let numericalFeatures = dataFeatures.map{ row in numericalColumns.map{ row[$0] } }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Categorize Categorical Features with Ordinal values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "var categoricalValues = Array(repeating: Set<Int32>(), count: 2)\n",
    "\n",
    "for record in categoricalFeatures {\n",
    "    categoricalValues[0].insert(record[0])\n",
    "    categoricalValues[1].insert(record[1])\n",
    "}\n",
    "\n",
    "let sortedCategoricalValues = [categoricalValues[0].sorted(), categoricalValues[1].sorted()]\n",
    "\n",
    "let ordinalCategoricalFeatures = categoricalFeatures.map{ [Int32(sortedCategoricalValues[0].firstIndex(of:$0[0])!), \n",
    "                                                           Int32(sortedCategoricalValues[1].firstIndex(of:$0[1])!)] }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split Train and Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "let trainPercentage:Float = 0.8\n",
    "let numTrainRecords = Int(ceil(Float(numRecords) * trainPercentage))\n",
    "let numTestRecords = numRecords - numTrainRecords\n",
    "\n",
    "func matrixTranspose<T>(_ matrix: [[T]]) -> [[T]] {\n",
    "    if matrix.isEmpty {return matrix}\n",
    "    var result = [[T]]()\n",
    "    for index in 0..<matrix.first!.count {\n",
    "        result.append(matrix.map{$0[index]})\n",
    "    }\n",
    "    return result\n",
    "}\n",
    "\n",
    "let xCategoricalAllTrain = matrixTranspose(Array(ordinalCategoricalFeatures[0..<numTrainRecords]))\n",
    "let xCategoricalAllTest = matrixTranspose(Array(ordinalCategoricalFeatures[numTrainRecords...]))\n",
    "let xNumericalAllTrain = Array(Array(numericalFeatures[0..<numTrainRecords]).joined())\n",
    "let xNumericalAllTest = Array(Array(numericalFeatures[numTrainRecords...]).joined())\n",
    "let yAllTrain = Array(Array(dataLabels[0..<numTrainRecords]).joined())\n",
    "let yAllTest = Array(Array(dataLabels[numTrainRecords...]).joined())\n",
    "\n",
    "let XCategoricalTrain = xCategoricalAllTrain.enumerated().map{ (offset, element) in \n",
    "    Tensor<Int32>(element).reshaped(to: TensorShape([numTrainRecords, 1]))\n",
    "}\n",
    "let XCategoricalTest = xCategoricalAllTest.enumerated().map{ (offset, element) in \n",
    "    Tensor<Int32>(element).reshaped(to: TensorShape([numTestRecords, 1]))\n",
    "}\n",
    "\n",
    "let XNumericalTrainDeNorm = Tensor<Float>(xNumericalAllTrain).reshaped(to: TensorShape([numTrainRecords, numNumericalFeatures]))\n",
    "let XNumericalTestDeNorm = Tensor<Float>(xNumericalAllTest).reshaped(to: TensorShape([numTestRecords, numNumericalFeatures]))\n",
    "let YTrain = Tensor<Float>(yAllTrain).reshaped(to: TensorShape([numTrainRecords, numLabels]))\n",
    "let YTest = Tensor<Float>(yAllTest).reshaped(to: TensorShape([numTestRecords, numLabels]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Normalize Numerical Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 3.7323444,  11.655556,  11.091472, 0.55681896,   6.309071,   69.04397,  3.7524028,  406.80493,\r\n",
      "   18.402004,  356.88275,  12.641355]] [[  8.972491,  23.946005,   6.836427, 0.11637194, 0.69087964,   27.80315,   2.085469,  167.87407,\r\n",
      "   2.1604173,   90.81331,  7.0751004]]\r\n"
     ]
    }
   ],
   "source": [
    "let mean = XNumericalTrainDeNorm.mean(alongAxes: 0)\n",
    "let std = XNumericalTrainDeNorm.standardDeviation(alongAxes: 0)\n",
    "\n",
    "print(mean, std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "let XNumericalTrain = (XNumericalTrainDeNorm - mean)/std\n",
    "let XNumericalTest = (XNumericalTestDeNorm - mean)/std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training shapes [405, 11] [405, 1] [405, 1] [405, 1]\r\n",
      "Testing shapes  [101, 11] [101, 1] [101, 1] [101, 1]\r\n"
     ]
    }
   ],
   "source": [
    "print(\"Training shapes \\(XNumericalTrain.shape) \\(XCategoricalTrain[0].shape) \\(XCategoricalTrain[1].shape) \\(YTrain.shape)\")\n",
    "print(\"Testing shapes  \\(XNumericalTest.shape) \\(XCategoricalTest[0].shape) \\(XCategoricalTest[1].shape) \\(YTest.shape)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HDJZCCgqohIC"
   },
   "outputs": [],
   "source": [
    "struct MultiInputs<N: Differentiable, C>: Differentiable {\n",
    "  var numerical: N\n",
    "  \n",
    "  @noDerivative\n",
    "  var categorical: C\n",
    "\n",
    "  @differentiable\n",
    "  init(numerical: N, categorical: C) {\n",
    "    self.numerical = numerical\n",
    "    self.categorical = categorical\n",
    "  }\n",
    "}\n",
    "\n",
    "struct RegressionModel: Module {\n",
    "    var embedding1 = TensorFlow.Embedding<Float>(vocabularySize: 2, embeddingSize: 2)\n",
    "    var embedding2 = TensorFlow.Embedding<Float>(vocabularySize: 9, embeddingSize: 5)\n",
    "    var allInputConcatLayer = Dense<Float>(inputSize: (11 + 2 + 5), outputSize: 64, activation: relu)\n",
    "    var hiddenLayer = Dense<Float>(inputSize: 64, outputSize: 32, activation: relu)\n",
    "    var outputLayer = Dense<Float>(inputSize: 32, outputSize: 1)\n",
    "    \n",
    "    @differentiable\n",
    "    func callAsFunction(_ input: MultiInputs<[Tensor<Float>], [Tensor<Int32>]>) -> Tensor<Float> {\n",
    "        let embeddingOutput1 = embedding1(input.categorical[0])\n",
    "        let embeddingOutput1Reshaped = embeddingOutput1.reshaped(to: \n",
    "            TensorShape([embeddingOutput1.shape[0], embeddingOutput1.shape[2]]))\n",
    "        let embeddingOutput2 = embedding2(input.categorical[1])\n",
    "        let embeddingOutput2Reshaped = embeddingOutput2.reshaped(to: \n",
    "            TensorShape([embeddingOutput2.shape[0], embeddingOutput2.shape[2]]))\n",
    "        let allConcat = Tensor<Float>(concatenating: [input.numerical[0], embeddingOutput1Reshaped, embeddingOutput2Reshaped], alongAxis: 1)\n",
    "        return allConcat.sequenced(through: allInputConcatLayer, hiddenLayer, outputLayer)\n",
    "    }\n",
    "}\n",
    "\n",
    "var model = RegressionModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JK0Vj7bSohIF"
   },
   "outputs": [],
   "source": [
    "let optimizer = RMSProp(for: model, learningRate: 0.001)\n",
    "Context.local.learningPhase = .training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gY8C7yHJohIH"
   },
   "outputs": [],
   "source": [
    "let epochCount = 500\n",
    "let batchSize = 32\n",
    "let numberOfBatch = Int(ceil(Double(numTrainRecords) / Double(batchSize)))\n",
    "let shuffle = true\n",
    "\n",
    "func mae(predictions: Tensor<Float>, truths: Tensor<Float>) -> Float {\n",
    "    return abs(Tensor<Float>(predictions - truths)).mean().scalarized()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 203
    },
    "colab_type": "code",
    "id": "L9bU9HsdohIK",
    "outputId": "692b81c5-3286-4e7c-9246-eb56e6a3eaee"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: MSE: 563.8485, MAE: 21.897787\n",
      "Epoch 2: MSE: 490.1869, MAE: 20.34883\n",
      "Epoch 3: MSE: 409.94077, MAE: 18.485403\n",
      "Epoch 4: MSE: 318.39008, MAE: 16.091515\n",
      "Epoch 5: MSE: 226.24895, MAE: 13.2135315\n",
      "Epoch 6: MSE: 145.42302, MAE: 10.061981\n",
      "Epoch 7: MSE: 87.54926, MAE: 7.1168065\n",
      "Epoch 8: MSE: 54.246696, MAE: 5.061981\n",
      "Epoch 9: MSE: 37.383755, MAE: 3.9668653\n",
      "Epoch 10: MSE: 30.753527, MAE: 3.5773091\n",
      "Epoch 11: MSE: 27.649105, MAE: 3.4347148\n",
      "Epoch 12: MSE: 25.568367, MAE: 3.3387656\n",
      "Epoch 13: MSE: 24.153414, MAE: 3.2524211\n",
      "Epoch 14: MSE: 23.258257, MAE: 3.1831157\n",
      "Epoch 15: MSE: 22.236786, MAE: 3.086512\n",
      "Epoch 16: MSE: 21.149607, MAE: 3.0255835\n",
      "Epoch 17: MSE: 20.517857, MAE: 2.9345677\n",
      "Epoch 18: MSE: 19.540665, MAE: 2.8829193\n",
      "Epoch 19: MSE: 19.098837, MAE: 2.8591158\n",
      "Epoch 20: MSE: 18.32507, MAE: 2.7741556\n",
      "Epoch 21: MSE: 18.030825, MAE: 2.7426963\n",
      "Epoch 22: MSE: 17.133114, MAE: 2.6880503\n",
      "Epoch 23: MSE: 16.672474, MAE: 2.639018\n",
      "Epoch 24: MSE: 16.25662, MAE: 2.6031988\n",
      "Epoch 25: MSE: 15.577813, MAE: 2.5601408\n",
      "Epoch 26: MSE: 15.305609, MAE: 2.4771209\n",
      "Epoch 27: MSE: 14.704122, MAE: 2.4390204\n",
      "Epoch 28: MSE: 14.267274, MAE: 2.399069\n",
      "Epoch 29: MSE: 13.991804, MAE: 2.3293169\n",
      "Epoch 30: MSE: 13.639078, MAE: 2.3358786\n",
      "Epoch 31: MSE: 13.256102, MAE: 2.2952425\n",
      "Epoch 32: MSE: 12.969169, MAE: 2.2483637\n",
      "Epoch 33: MSE: 12.715768, MAE: 2.2260659\n",
      "Epoch 34: MSE: 12.541288, MAE: 2.20436\n",
      "Epoch 35: MSE: 12.149067, MAE: 2.1638758\n",
      "Epoch 36: MSE: 11.965489, MAE: 2.1739528\n",
      "Epoch 37: MSE: 11.666312, MAE: 2.1257803\n",
      "Epoch 38: MSE: 11.588872, MAE: 2.1153336\n",
      "Epoch 39: MSE: 11.342696, MAE: 2.11995\n",
      "Epoch 40: MSE: 11.224479, MAE: 2.0955882\n",
      "Epoch 41: MSE: 10.954921, MAE: 2.0839536\n",
      "Epoch 42: MSE: 10.72561, MAE: 2.0597806\n",
      "Epoch 43: MSE: 10.692925, MAE: 2.0510945\n",
      "Epoch 44: MSE: 10.491268, MAE: 2.0333962\n",
      "Epoch 45: MSE: 10.213952, MAE: 2.0255418\n",
      "Epoch 46: MSE: 10.163985, MAE: 2.0081496\n",
      "Epoch 47: MSE: 9.973029, MAE: 1.9843687\n",
      "Epoch 48: MSE: 9.833805, MAE: 2.0151427\n",
      "Epoch 49: MSE: 9.785643, MAE: 1.9844519\n",
      "Epoch 50: MSE: 9.692669, MAE: 1.9747437\n",
      "Epoch 51: MSE: 9.59797, MAE: 1.9685735\n",
      "Epoch 52: MSE: 9.414652, MAE: 1.9557978\n",
      "Epoch 53: MSE: 9.346058, MAE: 1.9640431\n",
      "Epoch 54: MSE: 9.324228, MAE: 1.9312731\n",
      "Epoch 55: MSE: 9.13505, MAE: 1.9366646\n",
      "Epoch 56: MSE: 9.021999, MAE: 1.9395007\n",
      "Epoch 57: MSE: 8.855015, MAE: 1.9247246\n",
      "Epoch 58: MSE: 8.820449, MAE: 1.9121768\n",
      "Epoch 59: MSE: 8.7876, MAE: 1.908528\n",
      "Epoch 60: MSE: 8.492888, MAE: 1.8950193\n",
      "Epoch 61: MSE: 8.597911, MAE: 1.8926057\n",
      "Epoch 62: MSE: 8.527984, MAE: 1.8770404\n",
      "Epoch 63: MSE: 8.284084, MAE: 1.8637074\n",
      "Epoch 64: MSE: 8.463932, MAE: 1.8670218\n",
      "Epoch 65: MSE: 8.173957, MAE: 1.8451774\n",
      "Epoch 66: MSE: 8.195211, MAE: 1.8639177\n",
      "Epoch 67: MSE: 8.086954, MAE: 1.8650143\n",
      "Epoch 68: MSE: 8.082907, MAE: 1.84172\n",
      "Epoch 69: MSE: 8.01417, MAE: 1.8439617\n",
      "Epoch 70: MSE: 7.958613, MAE: 1.8371694\n",
      "Epoch 71: MSE: 7.766507, MAE: 1.8419429\n",
      "Epoch 72: MSE: 7.7833047, MAE: 1.8103123\n",
      "Epoch 73: MSE: 7.745486, MAE: 1.8208592\n",
      "Epoch 74: MSE: 7.481038, MAE: 1.8059318\n",
      "Epoch 75: MSE: 7.5969005, MAE: 1.8083751\n",
      "Epoch 76: MSE: 7.4599156, MAE: 1.7793483\n",
      "Epoch 77: MSE: 7.553805, MAE: 1.787345\n",
      "Epoch 78: MSE: 7.4719954, MAE: 1.7764599\n",
      "Epoch 79: MSE: 7.0357065, MAE: 1.7789843\n",
      "Epoch 80: MSE: 7.3746243, MAE: 1.7632285\n",
      "Epoch 81: MSE: 7.1351414, MAE: 1.7598977\n",
      "Epoch 82: MSE: 7.224724, MAE: 1.7577161\n",
      "Epoch 83: MSE: 7.199143, MAE: 1.7609011\n",
      "Epoch 84: MSE: 7.0020065, MAE: 1.7503561\n",
      "Epoch 85: MSE: 7.0680637, MAE: 1.7474208\n",
      "Epoch 86: MSE: 6.900363, MAE: 1.7447444\n",
      "Epoch 87: MSE: 6.8122683, MAE: 1.7240546\n",
      "Epoch 88: MSE: 6.816235, MAE: 1.7262473\n",
      "Epoch 89: MSE: 6.661323, MAE: 1.7293917\n",
      "Epoch 90: MSE: 6.821665, MAE: 1.7273891\n",
      "Epoch 91: MSE: 6.638402, MAE: 1.7305195\n",
      "Epoch 92: MSE: 6.6883187, MAE: 1.6999588\n",
      "Epoch 93: MSE: 6.507513, MAE: 1.6876805\n",
      "Epoch 94: MSE: 6.565255, MAE: 1.6916703\n",
      "Epoch 95: MSE: 6.4810743, MAE: 1.6851735\n",
      "Epoch 96: MSE: 6.4107485, MAE: 1.6871611\n",
      "Epoch 97: MSE: 6.3566027, MAE: 1.6719252\n",
      "Epoch 98: MSE: 6.208868, MAE: 1.6610456\n",
      "Epoch 99: MSE: 6.2077837, MAE: 1.651793\n",
      "Epoch 100: MSE: 6.2100887, MAE: 1.6417753\n",
      "Epoch 101: MSE: 6.0738926, MAE: 1.6516392\n",
      "Epoch 102: MSE: 5.946229, MAE: 1.634608\n",
      "Epoch 103: MSE: 6.0404515, MAE: 1.6267707\n",
      "Epoch 104: MSE: 6.139101, MAE: 1.6230447\n",
      "Epoch 105: MSE: 5.9049635, MAE: 1.625705\n",
      "Epoch 106: MSE: 5.8598523, MAE: 1.6175975\n",
      "Epoch 107: MSE: 5.893387, MAE: 1.6009096\n",
      "Epoch 108: MSE: 5.6666293, MAE: 1.5951492\n",
      "Epoch 109: MSE: 5.8128324, MAE: 1.5998254\n",
      "Epoch 110: MSE: 5.632924, MAE: 1.5815759\n",
      "Epoch 111: MSE: 5.7637362, MAE: 1.5895613\n",
      "Epoch 112: MSE: 5.6500726, MAE: 1.5837345\n",
      "Epoch 113: MSE: 5.5668902, MAE: 1.5724136\n",
      "Epoch 114: MSE: 5.611633, MAE: 1.5679017\n",
      "Epoch 115: MSE: 5.66127, MAE: 1.5581619\n",
      "Epoch 116: MSE: 5.4855413, MAE: 1.5685347\n",
      "Epoch 117: MSE: 5.381317, MAE: 1.5452311\n",
      "Epoch 118: MSE: 5.4263515, MAE: 1.5408094\n",
      "Epoch 119: MSE: 5.419095, MAE: 1.5381761\n",
      "Epoch 120: MSE: 5.196413, MAE: 1.5256068\n",
      "Epoch 121: MSE: 5.2355294, MAE: 1.5190406\n",
      "Epoch 122: MSE: 5.3632474, MAE: 1.52685\n",
      "Epoch 123: MSE: 5.3699045, MAE: 1.5233344\n",
      "Epoch 124: MSE: 5.179612, MAE: 1.5218498\n",
      "Epoch 125: MSE: 5.180559, MAE: 1.5148358\n",
      "Epoch 126: MSE: 5.112328, MAE: 1.5004134\n",
      "Epoch 127: MSE: 5.0085263, MAE: 1.4856452\n",
      "Epoch 128: MSE: 4.980833, MAE: 1.4889851\n",
      "Epoch 129: MSE: 4.9427195, MAE: 1.4732287\n",
      "Epoch 130: MSE: 5.119105, MAE: 1.4841772\n",
      "Epoch 131: MSE: 4.885818, MAE: 1.4904994\n",
      "Epoch 132: MSE: 4.988693, MAE: 1.4593352\n",
      "Epoch 133: MSE: 4.871461, MAE: 1.4717407\n",
      "Epoch 134: MSE: 4.8373957, MAE: 1.4586583\n",
      "Epoch 135: MSE: 4.8240237, MAE: 1.4628441\n",
      "Epoch 136: MSE: 4.8220816, MAE: 1.4463855\n",
      "Epoch 137: MSE: 4.8583, MAE: 1.4457456\n",
      "Epoch 138: MSE: 4.7475696, MAE: 1.4300927\n",
      "Epoch 139: MSE: 4.7553167, MAE: 1.4421718\n",
      "Epoch 140: MSE: 4.7234774, MAE: 1.4334056\n",
      "Epoch 141: MSE: 4.6565485, MAE: 1.4315807\n",
      "Epoch 142: MSE: 4.6947637, MAE: 1.4170299\n",
      "Epoch 143: MSE: 4.651696, MAE: 1.4375974\n",
      "Epoch 144: MSE: 4.52176, MAE: 1.4061731\n",
      "Epoch 145: MSE: 4.5853915, MAE: 1.4019585\n",
      "Epoch 146: MSE: 4.504251, MAE: 1.4025122\n",
      "Epoch 147: MSE: 4.5074096, MAE: 1.406161\n",
      "Epoch 148: MSE: 4.4576244, MAE: 1.3977159\n",
      "Epoch 149: MSE: 4.4182844, MAE: 1.3923738\n",
      "Epoch 150: MSE: 4.3966603, MAE: 1.3847296\n",
      "Epoch 151: MSE: 4.3536153, MAE: 1.3917459\n",
      "Epoch 152: MSE: 4.455376, MAE: 1.3730369\n",
      "Epoch 153: MSE: 4.2962384, MAE: 1.3737124\n",
      "Epoch 154: MSE: 4.40004, MAE: 1.369904\n",
      "Epoch 155: MSE: 4.3215213, MAE: 1.3657842\n",
      "Epoch 156: MSE: 4.3064213, MAE: 1.3662015\n",
      "Epoch 157: MSE: 4.2229004, MAE: 1.3463154\n",
      "Epoch 158: MSE: 4.179804, MAE: 1.3525329\n",
      "Epoch 159: MSE: 4.245711, MAE: 1.3496522\n",
      "Epoch 160: MSE: 4.0646544, MAE: 1.3447526\n",
      "Epoch 161: MSE: 4.149688, MAE: 1.3352822\n",
      "Epoch 162: MSE: 4.1467943, MAE: 1.333917\n",
      "Epoch 163: MSE: 4.002254, MAE: 1.3303877\n",
      "Epoch 164: MSE: 4.1002965, MAE: 1.323309\n",
      "Epoch 165: MSE: 4.091444, MAE: 1.3327113\n",
      "Epoch 166: MSE: 3.8296547, MAE: 1.3295038\n",
      "Epoch 167: MSE: 4.0972505, MAE: 1.313612\n",
      "Epoch 168: MSE: 3.8852816, MAE: 1.3232007\n",
      "Epoch 169: MSE: 3.926562, MAE: 1.2995985\n",
      "Epoch 170: MSE: 3.9603817, MAE: 1.3029506\n",
      "Epoch 171: MSE: 3.922916, MAE: 1.311976\n",
      "Epoch 172: MSE: 3.8455381, MAE: 1.2965373\n",
      "Epoch 173: MSE: 3.9406323, MAE: 1.2987425\n",
      "Epoch 174: MSE: 3.7066765, MAE: 1.2963477\n",
      "Epoch 175: MSE: 3.8024316, MAE: 1.2796477\n",
      "Epoch 176: MSE: 3.7963386, MAE: 1.2816443\n",
      "Epoch 177: MSE: 3.7094655, MAE: 1.274167\n",
      "Epoch 178: MSE: 3.8016703, MAE: 1.269392\n",
      "Epoch 179: MSE: 3.7384439, MAE: 1.2797109\n",
      "Epoch 180: MSE: 3.6355207, MAE: 1.2711742\n",
      "Epoch 181: MSE: 3.7238688, MAE: 1.2682843\n",
      "Epoch 182: MSE: 3.6098828, MAE: 1.2605057\n",
      "Epoch 183: MSE: 3.6071217, MAE: 1.2583817\n",
      "Epoch 184: MSE: 3.6361563, MAE: 1.2559798\n",
      "Epoch 185: MSE: 3.6335206, MAE: 1.2602084\n",
      "Epoch 186: MSE: 3.5620458, MAE: 1.2434545\n",
      "Epoch 187: MSE: 3.5417445, MAE: 1.2474364\n",
      "Epoch 188: MSE: 3.5028713, MAE: 1.2455022\n",
      "Epoch 189: MSE: 3.547349, MAE: 1.2428899\n",
      "Epoch 190: MSE: 3.515686, MAE: 1.2196186\n",
      "Epoch 191: MSE: 3.487761, MAE: 1.2366632\n",
      "Epoch 192: MSE: 3.3789585, MAE: 1.2281138\n",
      "Epoch 193: MSE: 3.3675642, MAE: 1.2160649\n",
      "Epoch 194: MSE: 3.4403017, MAE: 1.2155856\n",
      "Epoch 195: MSE: 3.384698, MAE: 1.2215449\n",
      "Epoch 196: MSE: 3.3874714, MAE: 1.2231401\n",
      "Epoch 197: MSE: 3.3304586, MAE: 1.202765\n",
      "Epoch 198: MSE: 3.318876, MAE: 1.2072877\n",
      "Epoch 199: MSE: 3.4457211, MAE: 1.2213657\n",
      "Epoch 200: MSE: 3.2489824, MAE: 1.2014811\n",
      "Epoch 201: MSE: 3.2637875, MAE: 1.1995883\n",
      "Epoch 202: MSE: 3.4246674, MAE: 1.2151371\n",
      "Epoch 203: MSE: 3.213217, MAE: 1.195847\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 204: MSE: 3.2787564, MAE: 1.1949987\n",
      "Epoch 205: MSE: 3.193905, MAE: 1.1847005\n",
      "Epoch 206: MSE: 3.2600913, MAE: 1.1892202\n",
      "Epoch 207: MSE: 3.2797556, MAE: 1.1785502\n",
      "Epoch 208: MSE: 3.149941, MAE: 1.1760744\n",
      "Epoch 209: MSE: 3.1763976, MAE: 1.1737677\n",
      "Epoch 210: MSE: 3.0601056, MAE: 1.1673385\n",
      "Epoch 211: MSE: 3.1165204, MAE: 1.1579679\n",
      "Epoch 212: MSE: 3.189201, MAE: 1.1653835\n",
      "Epoch 213: MSE: 3.0799482, MAE: 1.1646851\n",
      "Epoch 214: MSE: 3.0932424, MAE: 1.1592627\n",
      "Epoch 215: MSE: 3.1427445, MAE: 1.1671541\n",
      "Epoch 216: MSE: 3.0114882, MAE: 1.1532263\n",
      "Epoch 217: MSE: 3.0194883, MAE: 1.1518983\n",
      "Epoch 218: MSE: 3.00341, MAE: 1.1462104\n",
      "Epoch 219: MSE: 2.9979973, MAE: 1.1430625\n",
      "Epoch 220: MSE: 3.0376482, MAE: 1.1508343\n",
      "Epoch 221: MSE: 3.0058577, MAE: 1.1416032\n",
      "Epoch 222: MSE: 2.9216533, MAE: 1.131959\n",
      "Epoch 223: MSE: 2.9953341, MAE: 1.1414534\n",
      "Epoch 224: MSE: 2.950961, MAE: 1.1334188\n",
      "Epoch 225: MSE: 3.0177274, MAE: 1.1304051\n",
      "Epoch 226: MSE: 2.9203074, MAE: 1.1276654\n",
      "Epoch 227: MSE: 2.8763418, MAE: 1.1249375\n",
      "Epoch 228: MSE: 2.8406549, MAE: 1.1225954\n",
      "Epoch 229: MSE: 2.9261856, MAE: 1.1190126\n",
      "Epoch 230: MSE: 2.924388, MAE: 1.1206946\n",
      "Epoch 231: MSE: 2.8035855, MAE: 1.1069934\n",
      "Epoch 232: MSE: 2.829539, MAE: 1.1102431\n",
      "Epoch 233: MSE: 2.8335333, MAE: 1.1153126\n",
      "Epoch 234: MSE: 2.723364, MAE: 1.0981997\n",
      "Epoch 235: MSE: 2.871431, MAE: 1.0909985\n",
      "Epoch 236: MSE: 2.7080545, MAE: 1.1051035\n",
      "Epoch 237: MSE: 2.6718683, MAE: 1.0829041\n",
      "Epoch 238: MSE: 2.7950919, MAE: 1.0931844\n",
      "Epoch 239: MSE: 2.6870346, MAE: 1.0990925\n",
      "Epoch 240: MSE: 2.646422, MAE: 1.0826106\n",
      "Epoch 241: MSE: 2.720727, MAE: 1.0862575\n",
      "Epoch 242: MSE: 2.7416937, MAE: 1.0734508\n",
      "Epoch 243: MSE: 2.7767885, MAE: 1.0798959\n",
      "Epoch 244: MSE: 2.6324973, MAE: 1.0863367\n",
      "Epoch 245: MSE: 2.680981, MAE: 1.0773422\n",
      "Epoch 246: MSE: 2.7446897, MAE: 1.0765505\n",
      "Epoch 247: MSE: 2.6433916, MAE: 1.073416\n",
      "Epoch 248: MSE: 2.657889, MAE: 1.0658803\n",
      "Epoch 249: MSE: 2.5935106, MAE: 1.0693743\n",
      "Epoch 250: MSE: 2.6056073, MAE: 1.0595227\n",
      "Epoch 251: MSE: 2.6053643, MAE: 1.0663854\n",
      "Epoch 252: MSE: 2.524062, MAE: 1.0535862\n",
      "Epoch 253: MSE: 2.6123838, MAE: 1.0648139\n",
      "Epoch 254: MSE: 2.5457675, MAE: 1.0544243\n",
      "Epoch 255: MSE: 2.5323355, MAE: 1.0451717\n",
      "Epoch 256: MSE: 2.5309718, MAE: 1.0579613\n",
      "Epoch 257: MSE: 2.5957665, MAE: 1.0406156\n",
      "Epoch 258: MSE: 2.565458, MAE: 1.0494374\n",
      "Epoch 259: MSE: 2.4814935, MAE: 1.0412142\n",
      "Epoch 260: MSE: 2.409919, MAE: 1.0420406\n",
      "Epoch 261: MSE: 2.513194, MAE: 1.0408542\n",
      "Epoch 262: MSE: 2.4052873, MAE: 1.0245614\n",
      "Epoch 263: MSE: 2.4319656, MAE: 1.0368102\n",
      "Epoch 264: MSE: 2.4322135, MAE: 1.0211751\n",
      "Epoch 265: MSE: 2.4153836, MAE: 1.0249138\n",
      "Epoch 266: MSE: 2.513083, MAE: 1.0397445\n",
      "Epoch 267: MSE: 2.4324837, MAE: 1.0137092\n",
      "Epoch 268: MSE: 2.3625424, MAE: 1.0277674\n",
      "Epoch 269: MSE: 2.3464148, MAE: 1.0208663\n",
      "Epoch 270: MSE: 2.4271283, MAE: 1.0209746\n",
      "Epoch 271: MSE: 2.378082, MAE: 1.0083951\n",
      "Epoch 272: MSE: 2.3470926, MAE: 1.0193005\n",
      "Epoch 273: MSE: 2.4309118, MAE: 1.0181903\n",
      "Epoch 274: MSE: 2.314633, MAE: 0.9951162\n",
      "Epoch 275: MSE: 2.2377534, MAE: 0.9913879\n",
      "Epoch 276: MSE: 2.3397164, MAE: 1.0250504\n",
      "Epoch 277: MSE: 2.350208, MAE: 1.0057437\n",
      "Epoch 278: MSE: 2.2415965, MAE: 0.9928137\n",
      "Epoch 279: MSE: 2.3050396, MAE: 0.99444443\n",
      "Epoch 280: MSE: 2.3233557, MAE: 0.98689705\n",
      "Epoch 281: MSE: 2.2060614, MAE: 0.9799884\n",
      "Epoch 282: MSE: 2.1831858, MAE: 0.98127836\n",
      "Epoch 283: MSE: 2.1851902, MAE: 0.9870889\n",
      "Epoch 284: MSE: 2.2779105, MAE: 0.9909211\n",
      "Epoch 285: MSE: 2.254045, MAE: 0.9864549\n",
      "Epoch 286: MSE: 2.1865907, MAE: 0.9780172\n",
      "Epoch 287: MSE: 2.2738473, MAE: 0.98174167\n",
      "Epoch 288: MSE: 2.112638, MAE: 0.9531821\n",
      "Epoch 289: MSE: 2.1657138, MAE: 0.98107755\n",
      "Epoch 290: MSE: 2.238709, MAE: 0.9864407\n",
      "Epoch 291: MSE: 2.124693, MAE: 0.9574945\n",
      "Epoch 292: MSE: 2.1472528, MAE: 0.95975864\n",
      "Epoch 293: MSE: 2.1604757, MAE: 0.96683896\n",
      "Epoch 294: MSE: 2.1875918, MAE: 0.96340364\n",
      "Epoch 295: MSE: 2.0721202, MAE: 0.9603434\n",
      "Epoch 296: MSE: 2.0711699, MAE: 0.9588363\n",
      "Epoch 297: MSE: 2.072252, MAE: 0.95124555\n",
      "Epoch 298: MSE: 2.109758, MAE: 0.96924675\n",
      "Epoch 299: MSE: 2.1816623, MAE: 0.9487419\n",
      "Epoch 300: MSE: 2.087964, MAE: 0.9580799\n",
      "Epoch 301: MSE: 2.054572, MAE: 0.93964726\n",
      "Epoch 302: MSE: 2.0915027, MAE: 0.9411039\n",
      "Epoch 303: MSE: 2.0853255, MAE: 0.9534228\n",
      "Epoch 304: MSE: 2.0315783, MAE: 0.94771296\n",
      "Epoch 305: MSE: 2.0799203, MAE: 0.9350946\n",
      "Epoch 306: MSE: 2.028615, MAE: 0.9335805\n",
      "Epoch 307: MSE: 2.136452, MAE: 0.9407124\n",
      "Epoch 308: MSE: 2.0428584, MAE: 0.93293035\n",
      "Epoch 309: MSE: 1.9933469, MAE: 0.9179887\n",
      "Epoch 310: MSE: 1.914017, MAE: 0.9087535\n",
      "Epoch 311: MSE: 2.0305464, MAE: 0.9358241\n",
      "Epoch 312: MSE: 1.9294358, MAE: 0.9074316\n",
      "Epoch 313: MSE: 2.025413, MAE: 0.928121\n",
      "Epoch 314: MSE: 1.9903188, MAE: 0.9369837\n",
      "Epoch 315: MSE: 1.9398315, MAE: 0.9278947\n",
      "Epoch 316: MSE: 2.022138, MAE: 0.9346956\n",
      "Epoch 317: MSE: 1.8689067, MAE: 0.9114193\n",
      "Epoch 318: MSE: 1.9937809, MAE: 0.9112874\n",
      "Epoch 319: MSE: 1.9275603, MAE: 0.91441476\n",
      "Epoch 320: MSE: 1.9217099, MAE: 0.91687655\n",
      "Epoch 321: MSE: 1.8258101, MAE: 0.9010695\n",
      "Epoch 322: MSE: 1.9657152, MAE: 0.90893453\n",
      "Epoch 323: MSE: 1.8946882, MAE: 0.9144627\n",
      "Epoch 324: MSE: 1.8842689, MAE: 0.90335757\n",
      "Epoch 325: MSE: 1.923869, MAE: 0.9081718\n",
      "Epoch 326: MSE: 1.9027209, MAE: 0.90612805\n",
      "Epoch 327: MSE: 1.7171435, MAE: 0.87714905\n",
      "Epoch 328: MSE: 1.9228873, MAE: 0.8953664\n",
      "Epoch 329: MSE: 1.872956, MAE: 0.89673114\n",
      "Epoch 330: MSE: 1.8982686, MAE: 0.89726603\n",
      "Epoch 331: MSE: 1.8316065, MAE: 0.8799217\n",
      "Epoch 332: MSE: 1.8749511, MAE: 0.8949967\n",
      "Epoch 333: MSE: 1.7871933, MAE: 0.8658637\n",
      "Epoch 334: MSE: 1.8108534, MAE: 0.88674855\n",
      "Epoch 335: MSE: 1.8116276, MAE: 0.88643986\n",
      "Epoch 336: MSE: 1.7843156, MAE: 0.88518083\n",
      "Epoch 337: MSE: 1.7782519, MAE: 0.87818897\n",
      "Epoch 338: MSE: 1.696734, MAE: 0.8917195\n",
      "Epoch 339: MSE: 1.8287752, MAE: 0.85673034\n",
      "Epoch 340: MSE: 1.7347983, MAE: 0.88452256\n",
      "Epoch 341: MSE: 1.7518659, MAE: 0.87227446\n",
      "Epoch 342: MSE: 1.7936785, MAE: 0.8750442\n",
      "Epoch 343: MSE: 1.7508245, MAE: 0.8621444\n",
      "Epoch 344: MSE: 1.6901759, MAE: 0.85536504\n",
      "Epoch 345: MSE: 1.830828, MAE: 0.86224025\n",
      "Epoch 346: MSE: 1.7266076, MAE: 0.8618535\n",
      "Epoch 347: MSE: 1.7960165, MAE: 0.8556735\n",
      "Epoch 348: MSE: 1.7371742, MAE: 0.86679876\n",
      "Epoch 349: MSE: 1.7042847, MAE: 0.8455089\n",
      "Epoch 350: MSE: 1.7014121, MAE: 0.8650351\n",
      "Epoch 351: MSE: 1.7733631, MAE: 0.85128605\n",
      "Epoch 352: MSE: 1.6902397, MAE: 0.85362786\n",
      "Epoch 353: MSE: 1.5752406, MAE: 0.84779555\n",
      "Epoch 354: MSE: 1.6692097, MAE: 0.84654903\n",
      "Epoch 355: MSE: 1.7194207, MAE: 0.86772513\n",
      "Epoch 356: MSE: 1.687567, MAE: 0.85147274\n",
      "Epoch 357: MSE: 1.6846792, MAE: 0.8403253\n",
      "Epoch 358: MSE: 1.6094463, MAE: 0.81917584\n",
      "Epoch 359: MSE: 1.744499, MAE: 0.8569655\n",
      "Epoch 360: MSE: 1.597571, MAE: 0.8370486\n",
      "Epoch 361: MSE: 1.6000621, MAE: 0.8420865\n",
      "Epoch 362: MSE: 1.6803801, MAE: 0.85362446\n",
      "Epoch 363: MSE: 1.583533, MAE: 0.8149719\n",
      "Epoch 364: MSE: 1.677438, MAE: 0.8470214\n",
      "Epoch 365: MSE: 1.6141502, MAE: 0.8390893\n",
      "Epoch 366: MSE: 1.5767981, MAE: 0.8294172\n",
      "Epoch 367: MSE: 1.6480103, MAE: 0.83252245\n",
      "Epoch 368: MSE: 1.604374, MAE: 0.81613165\n",
      "Epoch 369: MSE: 1.5823607, MAE: 0.8344624\n",
      "Epoch 370: MSE: 1.5943716, MAE: 0.8161271\n",
      "Epoch 371: MSE: 1.560245, MAE: 0.8230972\n",
      "Epoch 372: MSE: 1.5746765, MAE: 0.8046683\n",
      "Epoch 373: MSE: 1.567232, MAE: 0.8101453\n",
      "Epoch 374: MSE: 1.5743917, MAE: 0.8331512\n",
      "Epoch 375: MSE: 1.5430043, MAE: 0.7987495\n",
      "Epoch 376: MSE: 1.5631648, MAE: 0.8117737\n",
      "Epoch 377: MSE: 1.5599557, MAE: 0.8047683\n",
      "Epoch 378: MSE: 1.5103018, MAE: 0.79126036\n",
      "Epoch 379: MSE: 1.5569285, MAE: 0.8164969\n",
      "Epoch 380: MSE: 1.5578727, MAE: 0.80516714\n",
      "Epoch 381: MSE: 1.5501163, MAE: 0.8106414\n",
      "Epoch 382: MSE: 1.5149595, MAE: 0.81481755\n",
      "Epoch 383: MSE: 1.5261294, MAE: 0.816616\n",
      "Epoch 384: MSE: 1.4624196, MAE: 0.77907366\n",
      "Epoch 385: MSE: 1.543644, MAE: 0.8135541\n",
      "Epoch 386: MSE: 1.4905117, MAE: 0.80163014\n",
      "Epoch 387: MSE: 1.4505284, MAE: 0.78914464\n",
      "Epoch 388: MSE: 1.4615631, MAE: 0.7837209\n",
      "Epoch 389: MSE: 1.5282891, MAE: 0.79747576\n",
      "Epoch 390: MSE: 1.3985691, MAE: 0.77458996\n",
      "Epoch 391: MSE: 1.4585178, MAE: 0.7699896\n",
      "Epoch 392: MSE: 1.3999965, MAE: 0.76995474\n",
      "Epoch 393: MSE: 1.4466729, MAE: 0.804534\n",
      "Epoch 394: MSE: 1.5202324, MAE: 0.7811641\n",
      "Epoch 395: MSE: 1.4426848, MAE: 0.779556\n",
      "Epoch 396: MSE: 1.4115843, MAE: 0.79038376\n",
      "Epoch 397: MSE: 1.4695401, MAE: 0.7799836\n",
      "Epoch 398: MSE: 1.4812719, MAE: 0.773244\n",
      "Epoch 399: MSE: 1.4271867, MAE: 0.7808864\n",
      "Epoch 400: MSE: 1.4606664, MAE: 0.78494513\n",
      "Epoch 401: MSE: 1.3858552, MAE: 0.7795451\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 402: MSE: 1.4210702, MAE: 0.777464\n",
      "Epoch 403: MSE: 1.4483842, MAE: 0.79974055\n",
      "Epoch 404: MSE: 1.369652, MAE: 0.7648208\n",
      "Epoch 405: MSE: 1.4732833, MAE: 0.7656081\n",
      "Epoch 406: MSE: 1.4478816, MAE: 0.77046955\n",
      "Epoch 407: MSE: 1.3706603, MAE: 0.7591768\n",
      "Epoch 408: MSE: 1.4402331, MAE: 0.7918941\n",
      "Epoch 409: MSE: 1.3116634, MAE: 0.7392591\n",
      "Epoch 410: MSE: 1.4156004, MAE: 0.7966983\n",
      "Epoch 411: MSE: 1.3975389, MAE: 0.765698\n",
      "Epoch 412: MSE: 1.3045399, MAE: 0.7416798\n",
      "Epoch 413: MSE: 1.3434751, MAE: 0.72998595\n",
      "Epoch 414: MSE: 1.4325271, MAE: 0.79817575\n",
      "Epoch 415: MSE: 1.4326291, MAE: 0.75326025\n",
      "Epoch 416: MSE: 1.4218795, MAE: 0.7653383\n",
      "Epoch 417: MSE: 1.3130962, MAE: 0.7489406\n",
      "Epoch 418: MSE: 1.423676, MAE: 0.76003736\n",
      "Epoch 419: MSE: 1.3367435, MAE: 0.74118805\n",
      "Epoch 420: MSE: 1.3991078, MAE: 0.75145906\n",
      "Epoch 421: MSE: 1.2954596, MAE: 0.73781836\n",
      "Epoch 422: MSE: 1.326693, MAE: 0.7477082\n",
      "Epoch 423: MSE: 1.3629708, MAE: 0.7441068\n",
      "Epoch 424: MSE: 1.3161595, MAE: 0.7519764\n",
      "Epoch 425: MSE: 1.3708088, MAE: 0.7604971\n",
      "Epoch 426: MSE: 1.3889312, MAE: 0.7285537\n",
      "Epoch 427: MSE: 1.1926128, MAE: 0.71312463\n",
      "Epoch 428: MSE: 1.288421, MAE: 0.7453292\n",
      "Epoch 429: MSE: 1.2989405, MAE: 0.74641347\n",
      "Epoch 430: MSE: 1.3186971, MAE: 0.7399155\n",
      "Epoch 431: MSE: 1.2093847, MAE: 0.73865765\n",
      "Epoch 432: MSE: 1.3712281, MAE: 0.7265248\n",
      "Epoch 433: MSE: 1.3039602, MAE: 0.72345424\n",
      "Epoch 434: MSE: 1.2382566, MAE: 0.7416929\n",
      "Epoch 435: MSE: 1.3437445, MAE: 0.7392472\n",
      "Epoch 436: MSE: 1.1725924, MAE: 0.7147363\n",
      "Epoch 437: MSE: 1.3527095, MAE: 0.7605342\n",
      "Epoch 438: MSE: 1.2578385, MAE: 0.72616476\n",
      "Epoch 439: MSE: 1.32308, MAE: 0.71510386\n",
      "Epoch 440: MSE: 1.2383435, MAE: 0.7225565\n",
      "Epoch 441: MSE: 1.2841351, MAE: 0.7464715\n",
      "Epoch 442: MSE: 1.2102201, MAE: 0.7060996\n",
      "Epoch 443: MSE: 1.3666624, MAE: 0.710858\n",
      "Epoch 444: MSE: 1.2971724, MAE: 0.7367272\n",
      "Epoch 445: MSE: 1.1724808, MAE: 0.7147944\n",
      "Epoch 446: MSE: 1.2451942, MAE: 0.7406063\n",
      "Epoch 447: MSE: 1.1641409, MAE: 0.6970237\n",
      "Epoch 448: MSE: 1.2953551, MAE: 0.72705686\n",
      "Epoch 449: MSE: 1.2373221, MAE: 0.69231576\n",
      "Epoch 450: MSE: 1.2636294, MAE: 0.729143\n",
      "Epoch 451: MSE: 1.2106988, MAE: 0.69838226\n",
      "Epoch 452: MSE: 1.2215049, MAE: 0.74099773\n",
      "Epoch 453: MSE: 1.2122213, MAE: 0.6931285\n",
      "Epoch 454: MSE: 1.1334904, MAE: 0.6891489\n",
      "Epoch 455: MSE: 1.166649, MAE: 0.7277373\n",
      "Epoch 456: MSE: 1.1676203, MAE: 0.6960738\n",
      "Epoch 457: MSE: 1.1943446, MAE: 0.7036345\n",
      "Epoch 458: MSE: 1.1643442, MAE: 0.7073822\n",
      "Epoch 459: MSE: 1.2559706, MAE: 0.7377521\n",
      "Epoch 460: MSE: 1.1915557, MAE: 0.70048183\n",
      "Epoch 461: MSE: 1.1031919, MAE: 0.7166188\n",
      "Epoch 462: MSE: 1.2162076, MAE: 0.7197466\n",
      "Epoch 463: MSE: 1.23118, MAE: 0.7114791\n",
      "Epoch 464: MSE: 1.0853324, MAE: 0.6826811\n",
      "Epoch 465: MSE: 1.138242, MAE: 0.6936029\n",
      "Epoch 466: MSE: 1.1886632, MAE: 0.7179916\n",
      "Epoch 467: MSE: 1.2075789, MAE: 0.70985603\n",
      "Epoch 468: MSE: 1.1911552, MAE: 0.72898775\n",
      "Epoch 469: MSE: 1.0533464, MAE: 0.6862782\n",
      "Epoch 470: MSE: 1.2711514, MAE: 0.69577974\n",
      "Epoch 471: MSE: 1.0838639, MAE: 0.65425926\n",
      "Epoch 472: MSE: 1.2004894, MAE: 0.7060438\n",
      "Epoch 473: MSE: 1.1704158, MAE: 0.68375134\n",
      "Epoch 474: MSE: 1.1074106, MAE: 0.66828126\n",
      "Epoch 475: MSE: 1.1958438, MAE: 0.71551603\n",
      "Epoch 476: MSE: 1.1128808, MAE: 0.6762175\n",
      "Epoch 477: MSE: 1.0746759, MAE: 0.7012682\n",
      "Epoch 478: MSE: 1.1574543, MAE: 0.6950093\n",
      "Epoch 479: MSE: 1.1034827, MAE: 0.68711954\n",
      "Epoch 480: MSE: 1.1949328, MAE: 0.6923952\n",
      "Epoch 481: MSE: 1.0852277, MAE: 0.66560644\n",
      "Epoch 482: MSE: 1.0476661, MAE: 0.6771834\n",
      "Epoch 483: MSE: 1.1738807, MAE: 0.69114304\n",
      "Epoch 484: MSE: 1.1702425, MAE: 0.70230526\n",
      "Epoch 485: MSE: 1.1070521, MAE: 0.68468225\n",
      "Epoch 486: MSE: 1.116772, MAE: 0.6614243\n",
      "Epoch 487: MSE: 1.0532563, MAE: 0.6864592\n",
      "Epoch 488: MSE: 1.1923449, MAE: 0.6747587\n",
      "Epoch 489: MSE: 1.0836462, MAE: 0.6703021\n",
      "Epoch 490: MSE: 1.1285284, MAE: 0.7010635\n",
      "Epoch 491: MSE: 1.0752164, MAE: 0.64147604\n",
      "Epoch 492: MSE: 1.1279563, MAE: 0.70792526\n",
      "Epoch 493: MSE: 1.0593505, MAE: 0.6657989\n",
      "Epoch 494: MSE: 1.0906881, MAE: 0.6982375\n",
      "Epoch 495: MSE: 1.034799, MAE: 0.6756476\n",
      "Epoch 496: MSE: 1.0946778, MAE: 0.68111193\n",
      "Epoch 497: MSE: 1.089342, MAE: 0.69773394\n",
      "Epoch 498: MSE: 1.1187088, MAE: 0.69519717\n",
      "Epoch 499: MSE: 1.0985531, MAE: 0.65445745\n",
      "Epoch 500: MSE: 1.0349859, MAE: 0.63037324\n"
     ]
    }
   ],
   "source": [
    "for epoch in 1...epochCount {\n",
    "    var epochLoss: Float = 0\n",
    "    var epochMAE: Float = 0\n",
    "    var batchCount: Int = 0\n",
    "    var batchArray = Array(repeating: false, count: numberOfBatch)\n",
    "    for batch in 0..<numberOfBatch {\n",
    "        var r = batch\n",
    "        if shuffle {\n",
    "            while true {\n",
    "                r = Int.random(in: 0..<numberOfBatch)\n",
    "                if !batchArray[r] {\n",
    "                    batchArray[r] = true\n",
    "                    break\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "        \n",
    "        let batchStart = r * batchSize\n",
    "        let batchEnd = min(numTrainRecords, batchStart + batchSize)\n",
    "        let (loss, grad) = model.valueWithGradient { (model: RegressionModel) -> Tensor<Float> in\n",
    "            let multiInput = MultiInputs(numerical: [XNumericalTrain[batchStart..<batchEnd]],\n",
    "                                         categorical: [XCategoricalTrain[0][batchStart..<batchEnd],\n",
    "                                                       XCategoricalTrain[1][batchStart..<batchEnd]])\n",
    "            let logits = model(multiInput)\n",
    "            return meanSquaredError(predicted: logits, expected: YTrain[batchStart..<batchEnd])\n",
    "        }\n",
    "        optimizer.update(&model, along: grad)\n",
    "        \n",
    "        let multiInput = MultiInputs(numerical: [XNumericalTrain[batchStart..<batchEnd]],\n",
    "                                     categorical: [XCategoricalTrain[0][batchStart..<batchEnd],\n",
    "                                                   XCategoricalTrain[1][batchStart..<batchEnd]])\n",
    "        let logits = model(multiInput)\n",
    "        epochMAE += mae(predictions: logits, truths: YTrain[batchStart..<batchEnd])\n",
    "        epochLoss += loss.scalarized()\n",
    "        batchCount += 1\n",
    "    }\n",
    "    epochMAE /= Float(batchCount)\n",
    "    epochLoss /= Float(batchCount)\n",
    "\n",
    "    print(\"Epoch \\(epoch): MSE: \\(epochLoss), MAE: \\(epochMAE)\")\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE: 0.12778322, MAE: 0.022995783\r\n"
     ]
    }
   ],
   "source": [
    "Context.local.learningPhase = .inference\n",
    "\n",
    "let multiInputTest = MultiInputs(numerical: [XNumericalTest],\n",
    "                                 categorical: [XCategoricalTest[0],\n",
    "                                               XCategoricalTest[1]])\n",
    "\n",
    "let prediction = model(multiInputTest)\n",
    "\n",
    "let predictionMse = meanSquaredError(predicted: prediction, expected: YTest).scalarized()/Float(numTestRecords)\n",
    "let predictionMae = mae(predictions: prediction, truths: YTest)/Float(numTestRecords)\n",
    "\n",
    "print(\"MSE: \\(predictionMse), MAE: \\(predictionMae)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Fkf29HLlohIP"
   },
   "source": [
    "# Export trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2, 2] [9, 5]\r\n"
     ]
    }
   ],
   "source": [
    "print(model.embedding1.embeddings.shape, model.embedding2.embeddings.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "let coremlModel = Model(version: 4,\n",
    "                        shortDescription: \"Regression\",\n",
    "                        author: \"Jacopo Mangiavacchi\",\n",
    "                        license: \"MIT\",\n",
    "                        userDefined: [\"SwiftCoremltoolsVersion\" : \"0.0.6\"]) {\n",
    "    Input(name: \"numericalInput\", shape: [11])\n",
    "    Input(name: \"categoricalInput1\", shape: [1])\n",
    "    Input(name: \"categoricalInput2\", shape: [1])\n",
    "    Output(name: \"output\", shape: [1])\n",
    "    NeuralNetwork {\n",
    "        Embedding(name: \"embedding1\",\n",
    "                     input: [\"categoricalInput1\"],\n",
    "                     output: [\"outEmbedding1\"],\n",
    "                     weight: model.embedding1.embeddings.transposed().flattened().scalars,\n",
    "                     inputDim: 2,\n",
    "                     outputChannels: 2)\n",
    "        Permute(name: \"permute1\",\n",
    "                     input: [\"outEmbedding1\"],\n",
    "                     output: [\"outPermute1\"],\n",
    "                     axis: [2, 1, 0, 3])\n",
    "        Flatten(name: \"flatten1\",\n",
    "                     input: [\"outPermute1\"],\n",
    "                     output: [\"outFlatten1\"],\n",
    "                     mode: .last)\n",
    "        Embedding(name: \"embedding2\",\n",
    "                     input: [\"categoricalInput2\"],\n",
    "                     output: [\"outEmbedding2\"],\n",
    "                     weight: model.embedding2.embeddings.transposed().flattened().scalars,\n",
    "                     inputDim: 9,\n",
    "                     outputChannels: 5)\n",
    "        Permute(name: \"permute2\",\n",
    "                     input: [\"outEmbedding2\"],\n",
    "                     output: [\"outPermute2\"],\n",
    "                     axis: [2, 1, 0, 3])\n",
    "        Flatten(name: \"flatten2\",\n",
    "                     input: [\"outPermute2\"],\n",
    "                     output: [\"outFlatten2\"],\n",
    "                     mode: .last)\n",
    "        Concat(name: \"concat\",\n",
    "                     input: [\"numericalInput\", \"outFlatten1\", \"outFlatten2\"],\n",
    "                     output: [\"outConcat\"])\n",
    "        InnerProduct(name: \"dense1\",\n",
    "                     input: [\"outConcat\"],\n",
    "                     output: [\"outDense1\"],\n",
    "                     weight: model.allInputConcatLayer.weight.transposed().flattened().scalars,\n",
    "                     bias: model.allInputConcatLayer.bias.flattened().scalars,\n",
    "                     inputChannels: 11 + 2 + 5,\n",
    "                     outputChannels: 64)\n",
    "        ReLu(name: \"Relu1\",\n",
    "             input: [\"outDense1\"],\n",
    "             output: [\"outRelu1\"])\n",
    "        InnerProduct(name: \"dense2\",\n",
    "                     input: [\"outRelu1\"],\n",
    "                     output: [\"outDense2\"],\n",
    "                     weight: model.hiddenLayer.weight.transposed().flattened().scalars,\n",
    "                     bias: model.hiddenLayer.bias.flattened().scalars,\n",
    "                     inputChannels: 64,\n",
    "                     outputChannels: 32)\n",
    "        ReLu(name: \"Relu2\",\n",
    "             input: [\"outDense2\"],\n",
    "             output: [\"outRelu2\"])\n",
    "        InnerProduct(name: \"dense3\",\n",
    "                     input: [\"outRelu2\"],\n",
    "                     output: [\"output\"],\n",
    "                     weight: model.outputLayer.weight.transposed().flattened().scalars,\n",
    "                     bias: model.outputLayer.bias.flattened().scalars,\n",
    "                     inputChannels: 32,\n",
    "                     outputChannels: 1)\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "let coreMLData = coremlModel.coreMLData\n",
    "try! coreMLData!.write(to: URL(fileURLWithPath: \"../model/s4tf_house_simplified_trained_model.mlmodel\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: [22.490442]\r\n",
      "MSE: 0.917915, MAE: 0.06598365\r\n"
     ]
    }
   ],
   "source": [
    "Context.local.learningPhase = .inference\n",
    "\n",
    "let numerical = [XNumericalTest[0].reshaped(to: TensorShape([1, 11]))]\n",
    "let cat1 = XCategoricalTest[0][0].reshaped(to: TensorShape([1, 1]))\n",
    "let cat2 = XCategoricalTest[1][0].reshaped(to: TensorShape([1, 1]))\n",
    "let categorical = [cat1, cat2]\n",
    "\n",
    "let multiInputTest = MultiInputs(numerical: numerical,\n",
    "                                 categorical: categorical)\n",
    "\n",
    "let prediction = model(multiInputTest)\n",
    "\n",
    "print(\"0: \\(prediction[0])\")\n",
    "\n",
    "let predictionMse = meanSquaredError(predicted: prediction, expected: YTest).scalarized()/Float(numTestRecords)\n",
    "let predictionMae = mae(predictions: prediction, truths: YTest)/Float(numTestRecords)\n",
    "\n",
    "print(\"MSE: \\(predictionMse), MAE: \\(predictionMae)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "func predictTestRecord(_ record: Int) {\n",
    "    Context.local.learningPhase = .inference\n",
    "\n",
    "    let numerical = [XNumericalTest[record].reshaped(to: TensorShape([1, 11]))]\n",
    "    let cat1 = XCategoricalTest[0][record].reshaped(to: TensorShape([1, 1]))\n",
    "    let cat2 = XCategoricalTest[1][record].reshaped(to: TensorShape([1, 1]))\n",
    "    let categorical = [cat1, cat2]\n",
    "\n",
    "    let multiInputTest = MultiInputs(numerical: numerical,\n",
    "                                     categorical: categorical)\n",
    "\n",
    "    let prediction = model(multiInputTest)\n",
    "\n",
    "    print(\"=== Test Record \\(record) ===\")\n",
    "    print(\"Numerical Input: \\(XNumericalTest[record])\")\n",
    "    print(\"Categorical Input: \\(XCategoricalTest[0][record]) \\(XCategoricalTest[1][record])\")\n",
    "    print(\"Prediction: \\(prediction[0])\")\n",
    "    print(\"Y: \\(YTest[record])\")\n",
    "    print(\"------------------\")\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Test Record 0 ===\r\n",
      "Numerical Input: [   -0.3748284,   -0.48674324,   -0.17428285,  -0.110154845,      0.373334,     0.6566175,\r\n",
      "   -0.07197559,   -0.61239314, -0.0009278871,       0.42733,    -0.4750964]\r\n",
      "Categorical Input: [0] [3]\r\n",
      "Prediction: [22.490442]\r\n",
      "Y: [23.8]\r\n",
      "------------------\r\n",
      "=== Test Record 17 ===\r\n",
      "Numerical Input: [  3.8583992, -0.48674324,   1.0251741,   1.1702225,   -1.239103,   1.1134001,   -1.085033,\r\n",
      "   1.5439851,    0.832245,  0.44065395,   2.5368752]\r\n",
      "Categorical Input: [0] [8]\r\n",
      "Prediction: [6.080847]\r\n",
      "Y: [5.0]\r\n",
      "------------------\r\n",
      "=== Test Record 87 ===\r\n",
      "Numerical Input: [-0.40829957, -0.48674324,  -1.2625705,  -0.5913707, -0.23892875, -0.24615799,  -0.5535939,\r\n",
      "  -1.2736031,  -0.2786522,  0.44065395, -0.45106846]\r\n",
      "Categorical Input: [0] [2]\r\n",
      "Prediction: [28.851]\r\n",
      "Y: [36.2]\r\n",
      "------------------\r\n"
     ]
    }
   ],
   "source": [
    "predictTestRecord(0)\n",
    "predictTestRecord(17)\n",
    "predictTestRecord(87)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Export trainable model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "let coremlModel = Model(version: 4,\n",
    "                        shortDescription: \"Regression\",\n",
    "                        author: \"Jacopo Mangiavacchi\",\n",
    "                        license: \"MIT\",\n",
    "                        userDefined: [\"SwiftCoremltoolsVersion\" : \"0.0.6\"]) {\n",
    "    Input(name: \"numericalInput\", shape: [11])\n",
    "    Input(name: \"categoricalInput1\", shape: [1])\n",
    "    Input(name: \"categoricalInput2\", shape: [1])\n",
    "    Output(name: \"output\", shape: [1])\n",
    "    TrainingInput(name: \"numericalInput\", shape: [11])\n",
    "    TrainingInput(name: \"categoricalInput1\", shape: [1])\n",
    "    TrainingInput(name: \"categoricalInput2\", shape: [1])\n",
    "    TrainingInput(name: \"output_true\", shape: [1])\n",
    "    NeuralNetwork(losses: [MSE(name: \"lossLayer\",\n",
    "                               input: \"output\",\n",
    "                               target: \"output_true\")],\n",
    "                  optimizer: SGD(learningRateDefault: 0.001,\n",
    "                                 learningRateMax: 0.3,\n",
    "                                 miniBatchSizeDefault: 32,\n",
    "                                 miniBatchSizeRange: [32],\n",
    "                                 momentumDefault: 0,\n",
    "                                 momentumMax: 1.0),\n",
    "                  epochDefault: 500,\n",
    "                  epochSet: [500],\n",
    "                  shuffle: true) {\n",
    "        Embedding(name: \"embedding1\",\n",
    "                     input: [\"categoricalInput1\"],\n",
    "                     output: [\"outEmbedding1\"],\n",
    "                     weight: model.embedding1.embeddings.transposed().flattened().scalars,\n",
    "                     inputDim: 2,\n",
    "                     outputChannels: 2)\n",
    "        Permute(name: \"permute1\",\n",
    "                     input: [\"outEmbedding1\"],\n",
    "                     output: [\"outPermute1\"],\n",
    "                     axis: [2, 1, 0, 3])\n",
    "        Flatten(name: \"flatten1\",\n",
    "                     input: [\"outPermute1\"],\n",
    "                     output: [\"outFlatten1\"],\n",
    "                     mode: .last)\n",
    "        Embedding(name: \"embedding2\",\n",
    "                     input: [\"categoricalInput2\"],\n",
    "                     output: [\"outEmbedding2\"],\n",
    "                     weight: model.embedding2.embeddings.transposed().flattened().scalars,\n",
    "                     inputDim: 9,\n",
    "                     outputChannels: 5)\n",
    "        Permute(name: \"permute2\",\n",
    "                     input: [\"outEmbedding2\"],\n",
    "                     output: [\"outPermute2\"],\n",
    "                     axis: [2, 1, 0, 3])\n",
    "        Flatten(name: \"flatten2\",\n",
    "                     input: [\"outPermute2\"],\n",
    "                     output: [\"outFlatten2\"],\n",
    "                     mode: .last)\n",
    "        Concat(name: \"concat\",\n",
    "                     input: [\"numericalInput\", \"outFlatten1\", \"outFlatten2\"],\n",
    "                     output: [\"outConcat\"])\n",
    "        InnerProduct(name: \"dense1\",\n",
    "                     input: [\"outConcat\"],\n",
    "                     output: [\"outDense1\"],\n",
    "                     inputChannels: 11 + 2 + 5,\n",
    "                     outputChannels: 64,\n",
    "                     updatable: true)\n",
    "        ReLu(name: \"Relu1\",\n",
    "             input: [\"outDense1\"],\n",
    "             output: [\"outRelu1\"])\n",
    "        InnerProduct(name: \"dense2\",\n",
    "                     input: [\"outRelu1\"],\n",
    "                     output: [\"outDense2\"],\n",
    "                     inputChannels: 64,\n",
    "                     outputChannels: 32,\n",
    "                     updatable: true)\n",
    "        ReLu(name: \"Relu2\",\n",
    "             input: [\"outDense2\"],\n",
    "             output: [\"outRelu2\"])\n",
    "        InnerProduct(name: \"dense3\",\n",
    "                     input: [\"outRelu2\"],\n",
    "                     output: [\"output\"],\n",
    "                     inputChannels: 32,\n",
    "                     outputChannels: 1,\n",
    "                     updatable: true)\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "let coreMLData = coremlModel.coreMLData\n",
    "try! coreMLData!.write(to: URL(fileURLWithPath: \"../model/s4tf_house_simplified_trainable_model.mlmodel\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "S4TF_House.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Swift",
   "language": "swift",
   "name": "swift"
  },
  "language_info": {
   "file_extension": ".swift",
   "mimetype": "text/x-swift",
   "name": "swift",
   "version": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
